import cv2
import os
from PIL import Image
import imagehash


def get_image_hash(image_path):
    """计算图片的哈希值"""
    with Image.open(image_path) as img:
        return imagehash.phash(img)  # 使用感知哈希值（pHash）


def delete_similar_images(folder_path, threshold=1):
    """删除文件夹下相似的图片"""
    image_hashes = {}
    deleted_count = 0
    for root, dirs, files in os.walk(folder_path):
        for file in files:
            if file.lower().endswith(('png', 'jpg', 'jpeg', 'gif', 'bmp')):
                file_path = os.path.join(root, file)
                hash_value = get_image_hash(file_path)

                # 检查是否已有相同哈希值的图片，若有则删除当前图片
                for stored_hash, stored_files in image_hashes.items():
                    if abs(hash_value - stored_hash) <= threshold:  # 根据阈值判断相似度
                        print(f"Deleting {file_path}, similar to {stored_files[0]}")
                        os.remove(file_path)
                        deleted_count += 1
                        break
                else:
                    # 如果没有找到相似的图片，则保存当前哈希值和文件
                    image_hashes.setdefault(hash_value, []).append(file_path)
    return deleted_count



def delete_similar_images2(folder_path, hash_size=8, threshold=5):
    """
    删除文件夹下相似的图片，保留最先出现的图片。

    参数:
    folder_path (str): 图片所在文件夹的路径。
    hash_size (int): 哈希值的大小，越大精度越高，计算越慢，默认8。
    threshold (int): 哈希值之间允许的最大差异，差异越小表示越相似，默认5。
    """
    # 获取文件夹中的所有图片文件
    image_files = [f for f in os.listdir(folder_path) if f.lower().endswith(('png', 'jpg', 'jpeg', 'bmp', 'gif'))]
    image_hashes = {}
    deleted_files = 0

    for image_file in image_files:
        image_path = os.path.join(folder_path, image_file)

        # 计算当前图片的哈希值
        try:
            img = Image.open(image_path)
            img_hash = imagehash.average_hash(img, hash_size=hash_size)
        except Exception as e:
            print(f"无法处理文件 {image_file}: {e}")
            continue

        # 检查当前哈希值与已保存的哈希值的差异
        found_similar = False
        for stored_hash, stored_image in image_hashes.items():
            if abs(img_hash - stored_hash) <= threshold:
                # 如果差异小于或等于阈值，则认为图片相似
                print(f"删除相似图片: {image_file}（与 {stored_image} 相似）")
                os.remove(image_path)
                deleted_files += 1
                found_similar = True
                break

        # 如果没有找到相似的图片，则保存哈希值
        if not found_similar:
            image_hashes[img_hash] = image_file
    return deleted_files


def extract_pics3(video_path, output_folder):
    """
    从给定的视频文件中提取帧，并将每一帧保存为图片到指定的输出文件夹。

    参数:
    video_path (str): 输入视频文件的路径。
    output_folder (str): 用于保存提取帧图片的输出文件夹路径。

    返回:
    int: 成功提取并保存的帧数。
    """
    # 创建VideoCapture对象来读取视频
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"无法打开视频文件: {video_path}")
        return 0

    # 检查输出文件夹是否存在，如果不存在则创建
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    frame_num = 0
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        # 构建保存图片的文件名，包括输出文件夹路径
        frame_name = os.path.join(output_folder, f"pic_{frame_num}.jpg")
        cv2.imwrite(frame_name, frame)
        frame_num += 1

    cap.release()
    return frame_num

def extract_pics2(video_path, output_folder, quality=90):
    """
    从给定的视频文件中提取帧，并将每一帧保存为压缩图片到指定的输出文件夹。

    参数:
    video_path (str): 输入视频文件的路径。
    output_folder (str): 用于保存提取帧图片的输出文件夹路径。
    quality (int): JPEG图片的压缩质量，范围为0到100，默认95。

    返回:
    int: 成功提取并保存的帧数。
    """
    # 创建VideoCapture对象来读取视频
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"无法打开视频文件: {video_path}")
        return 0

    # 检查输出文件夹是否存在，如果不存在则创建
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    frame_num = 0
    save_interval = 5  # 设置保存图片的间隔帧数，可调整
    saved_frames = 0  # 记录实际保存的帧数
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_num % save_interval == 0:
            # 构建保存图片的文件名，包括输出文件夹路径
            frame_name = os.path.join(output_folder, f"pic_{saved_frames}.jpg")
            # 使用JPEG压缩保存图片，quality参数控制压缩质量
            cv2.imwrite(frame_name, frame, [cv2.IMWRITE_JPEG_QUALITY, quality])
            saved_frames += 1
        frame_num += 1

    cap.release()
    return saved_frames

def extract_pics(video_path, output_folder):
    """
    从给定的视频文件中提取帧，并将每一帧保存为图片到指定的输出文件夹。

    参数:
    video_path (str): 输入视频文件的路径。
    output_folder (str): 用于保存提取帧图片的输出文件夹路径。

    返回:
    int: 成功提取并保存的帧数。
    """
    # 创建VideoCapture对象来读取视频
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"无法打开视频文件: {video_path}")
        return 0

    # 检查输出文件夹是否存在，如果不存在则创建
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    frame_num = 0
    save_interval = 5  # 设置保存图片的间隔帧数，可调整
    saved_frames = 0  # 记录实际保存的帧数
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if frame_num % save_interval == 0:
            # 构建保存图片的文件名，包括输出文件夹路径
            frame_name = os.path.join(output_folder, f"pic_{saved_frames}.jpg")
            cv2.imwrite(frame_name, frame)
            saved_frames += 1
        frame_num += 1

    cap.release()
    return saved_frames


if __name__ == "__main__":
    # 视频文件路径，替换为实际的视频文件路径
    video_path = "/Users/zengshan/Desktop/沧元图.S01E03.mp4"
    # 输出文件夹路径，用于保存提取的帧图片
    output_folder = "/Users/zengshan/Desktop/img"
    num_frames = extract_pics2(video_path, output_folder)
    print(f"共提取了{num_frames}帧图片并保存成功。")
    deleted_num = delete_similar_images2(output_folder,16,1)
    print(f"共删除了{deleted_num}张图片。")